#include "tools.h"
#include <ntimage.h>


ULONG GetThreadModeOffset()
{
	static int offset = 0;
	if (offset) return offset;

	UNICODE_STRING funcName = { 0 };
	RtlInitUnicodeString(&funcName, L"ExGetPreviousMode");
	PUCHAR func = (PUCHAR)MmGetSystemRoutineAddress(&funcName);
	if (func == NULL) return 0;
	PUCHAR temp = func;

	for (int i = 10; i < 50; i++)
	{
		if (temp[i] == 0xC3)
		{
			if (temp[i + 1] == 0x90 || temp[i + 1] == 0x0 || temp[i + 1] == 0xcc)
			{
				temp += i;
				break;
			}
		}
	}

	if (temp != func)
	{
		temp -= 4;
		offset = *(PULONG)temp;
	}

	return offset;
}

BOOLEAN SetThreadMode(PETHREAD thread, MODE newMode, MODE * oldMode)
{
	ULONG offset = GetThreadModeOffset();
	if (offset == 0) return FALSE;

	MODE mode = (MODE)*((PUCHAR)thread + offset);
	*((PUCHAR)thread + offset) = newMode;
	if (oldMode) *oldMode = mode;
	return TRUE;
}





NTSTATUS GetNtModuleBaseAndSize(ULONG64 * pModule, ULONG64 * pSize)
{
	if (pModule == NULL || pSize == NULL) return STATUS_UNSUCCESSFUL;
	ULONG BufferSize = PAGE_SIZE * 64;
	PVOID Buffer = ExAllocatePool(PagedPool, BufferSize);
	ULONG ReturnLength;
	NTSTATUS Status = ZwQuerySystemInformation(SystemModuleInformation,
		Buffer,
		BufferSize,
		&ReturnLength
		);

	if (Status == STATUS_INFO_LENGTH_MISMATCH) {
		ExFreePool(Buffer);
		return STATUS_INFO_LENGTH_MISMATCH;
	}

	PRTL_PROCESS_MODULES            Modules;
	PRTL_PROCESS_MODULE_INFORMATION ModuleInfo;
	Modules = (PRTL_PROCESS_MODULES)Buffer;
	ModuleInfo = &(Modules->Modules[0]);
	*pModule = ModuleInfo->ImageBase;
	*pSize = ModuleInfo->ImageSize;

	ExFreePool(Buffer);

	return Status;
}






ULONG GetWindowsVersionNumber()
{
	static ULONG gNumber = 0;
	if (gNumber != 0) return gNumber;

	RTL_OSVERSIONINFOW version = { 0 };
	RtlGetVersion(&version);

	if (version.dwMajorVersion <= 6) return 7;

	if (version.dwBuildNumber == 9600)
	{
		gNumber = 8;
	}
	else if (version.dwBuildNumber == 10240)
	{
		gNumber = 1507;
	}
	else if (version.dwBuildNumber == 10586)
	{
		gNumber = 1511;
	}
	else if (version.dwBuildNumber == 14393)
	{
		gNumber = 1607;
	}
	else if (version.dwBuildNumber == 15063)
	{
		gNumber = 1703;
	}
	else if (version.dwBuildNumber == 16299)
	{
		gNumber = 1709;
	}
	else if (version.dwBuildNumber == 17134)
	{
		gNumber = 1803;
	}
	else if (version.dwBuildNumber == 17763)
	{
		gNumber = 1809;
	}
	else if (version.dwBuildNumber == 18362)
	{
		gNumber = 1903;
	}
	else if (version.dwBuildNumber == 18363)
	{
		gNumber = 1909;
	}
	else if (version.dwBuildNumber == 19041)
	{
		gNumber = 2004;
	}
	else if (version.dwBuildNumber == 19042)
	{
		gNumber = 2009;
	}
	else if (version.dwBuildNumber == 19043)
	{
		gNumber = 2011;
	}
	else if (version.dwBuildNumber == 22200)
	{
		gNumber = 2012;
	}


	return gNumber;
}





LONG SafeSearchString(IN PUNICODE_STRING source, IN PUNICODE_STRING target, IN BOOLEAN CaseInSensitive)
{
	ASSERT(source != NULL && target != NULL);
	if (source == NULL || target == NULL || source->Buffer == NULL || target->Buffer == NULL)
		return STATUS_INVALID_PARAMETER;

	// Size mismatch
	if (source->Length < target->Length)
		return -1;

	USHORT diff = source->Length - target->Length;
	for (USHORT i = 0; i <= (diff / sizeof(WCHAR)); i++)
	{
		if (RtlCompareUnicodeStrings(
			source->Buffer + i,
			target->Length / sizeof(WCHAR),
			target->Buffer,
			target->Length / sizeof(WCHAR),
			CaseInSensitive
		) == 0)
		{
			return i;
		}
	}

	return -1;
}
